from fastapi import APIRouter, HTTPException
from model import model, processor
import scipy.io.wavfile
import numpy as np
from pydantic import BaseModel

router = APIRouter()

class T2STRequest(BaseModel):
    text: str
    src_lang: str
    tgt_lang: str

@router.post("/translate")
async def text_to_speech(request: T2STRequest):
    try:
        # Validate inputs
        if not request.text or not request.src_lang or not request.tgt_lang:
            raise HTTPException(status_code=400, detail="Text, source, and target languages are required")

        # Check if model and processor are loaded
        if model is None or processor is None:
            raise HTTPException(status_code=503, detail="Model not initialized")

        # Process text input
        text_inputs = processor(text=request.text, src_lang=request.src_lang, return_tensors="pt")

        # Generate speech
        audio_array = model.generate(**text_inputs, tgt_lang=request.tgt_lang)[0].cpu().numpy().squeeze()

        # Save output audio
        output_path = "output_t2st.wav"
        scipy.io.wavfile.write(output_path, rate=model.config.sampling_rate, data=audio_array)

        return {"message": "Text-to-speech translation completed", "output_file": output_path}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))